rm(list=ls())

# Data and model setup ----------------------------------------------------

load("survey_data.rda")
iter <- 3000

sum(exp_claassen_input[["x"]]) # sum respondents
unique(exp_claassen_input[["J"]]) # number countries
unique(exp_claassen_input[["K"]]) # number items

# Claassen's Model 5 ------------------------------------------------------

start_time <- Sys.time()

exp_claassen_m5 <- stan(file = 'supdem.stan.mod5.stan',
                        data = exp_claassen_input,
                        iter = iter,
                        seed = 0211,
                        chains= 4,
                        cores = 4,
                        thin = iter/500,
                        control = list(adapt_delta=0.99, stepsize=0.02, max_treedepth=11))

end_time <- Sys.time()
end_time - start_time

saveRDS(exp_claassen_m5, file = "claassen_mod5_output.rds")
exp_claassen_m5 <- readRDS(file = "claassen_mod5_output.rds")

# extract mean estimates

n.cntrys = length(unique(exp_claassen_input[["data"]][["country"]]))
cnt.names = as.character(sort(unique(exp_claassen_input[["data"]][["country"]])))

theta.out = rstan::extract(exp_claassen_m5, pars = c("theta"))[[1]]
theta.std = (theta.out - mean(as.vector(theta.out))) / sd(as.vector(theta.out)) # standardize
theta.out.t = apply(theta.std, 1, function(x) t(x) )
theta.out.df = data.frame(Country=rep(cnt.names, length.out=n.cntrys*33), 
                          Year=rep(1988:2020, each=n.cntrys), theta.out.t)
theta.pe = theta.out.df[,1:3]
theta.dim = dim(theta.out.df)[2]
theta.pe$SupDem_trim = apply(theta.out.df[,4:theta.dim], 1, mean)

first.yr = data.frame(Country=levels(as.factor(exp_claassen_input[["data"]]$country)),
                      First_yr = as.vector(by(exp_claassen_input[["data"]], exp_claassen_input[["data"]]$country, 
                                              function(x) min(as.numeric(x$year)))))

theta.pe = merge(theta.pe, first.yr, by="Country", all.x=TRUE)

theta.pe %>%
  mutate(SupDem_trim = case_when(First_yr <= Year ~ SupDem_trim,
                                 First_yr > Year ~ NA_real_)) %>%
  rename(citizen_support = SupDem_trim,
         country_name = Country,
         year = Year) %>%
  select(country_name, year, citizen_support) ->
  theta.pe2

saveRDS(theta.pe2, "claassen_mod5_mean_estimates.rds")

# extract estimates from each draw

theta.out.df = merge(theta.out.df, first.yr, by="Country", all.x=TRUE)

theta.pe2 %>% 
  left_join(theta.out.df, by = c("country_name" = "Country","year" = "Year")) %>% 
  select(-citizen_support) ->
  theta.pe2.draws

saveRDS(theta.pe2.draws, "claassen_mod5_draw_estimates.rds")

# Claassen's Model 6 ------------------------------------------------------

start_time <- Sys.time()

exp_claassen_m6 <- stan(file = 'supdem.stan.mod6.stan',
                        data = exp_claassen_input,
                        iter = iter,
                        seed = 0211,
                        chains= 4,
                        cores = 4,
                        thin = iter/500,
                        control = list(adapt_delta=0.99, stepsize=0.02, max_treedepth=11))

end_time <- Sys.time()
end_time - start_time

saveRDS(exp_claassen_m6, "claassen_mod6_mean_estimates.rds")